from datetime import datetime
import os
from typing import Dict, List, Optional, Generator, Tuple
import uuid
import json
from dotenv import load_dotenv
import requests
from flask import Response, jsonify, request

from ai_config import AIModelConfig

class AIChat:
    
    def __init__(self):
        self.active_dialogs: Dict[str, List[Dict]] = {}
        self.models: Dict[str, AIModelConfig] = {}
        self.register_models()
        self.presets = {
            "0": None,
            "1": {"role": "system", "content": "【语音模式】用口语化中文回复，3句话内"},
            "2": {"role": "system", "content": "【解题模式】分步骤解答计算机考研问题"},
            "3": {"role": "system", "content": "请生成一个408题目，题目要求如下：" +
                "1. 题目必须是计算机408考研相关的，不能是其他学科的题目。" +
                "2. 题目是选择题。" +
                "3. 题目内容必须是中文的，不能是英文的。" +
                "4. 题目内容和题目答案必须是完整的，不能是省略的。" +
                "5. 题目内容和题目答案必须是正确的，不能是错误的。"}
        }

    def register_models(self):
        """注册AI模型"""
        load_dotenv()
        self.models["DeepSeek-R1"] = AIModelConfig(
            name="DeepSeek-R1",api_base="http://10.2.8.77:3000/v1/chat/completions",
            api_key=os.getenv("DEEPSEEK_API_KEY") , model="DeepSeek-R1",)
        self.models["DeepSeek-V3"] = AIModelConfig(
            name="DeepSeek-V3",api_base="https://api.deepseek.com/chat/completions",
            api_key=os.getenv("DEEPSEEK_API_KEY_V3") , model="deepseek-chat",)
        self.models["doubao-thinking"] = AIModelConfig(
            name="doubao-thinking",api_base="https://ark.cn-beijing.volces.com/api/v3/chat/completions",
            api_key=os.getenv("ARK_API_KEY") , model="doubao-1-5-thinking-pro-m-250428",)
        self.models["doubao-thinking-vision"] = AIModelConfig(
            name="doubao-thinking-vision",api_base="https://ark.cn-beijing.volces.com/api/v3/chat/completions",
            api_key=os.getenv("ARK_API_KEY") , model="doubao-1-5-thinking-vision-pro-250428",)
        
        self.models["doubao-fine"] = AIModelConfig(
            name="doubao-fine",api_base="https://ark.cn-beijing.volces.com/api/v3/chat/completions",
            api_key=os.getenv("ARK_API_KEY") , model="ep-20250612183138-stmkt",)
        
        self.models["qwen-plus"] = AIModelConfig(
            name="qwen-plus",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model="qwen-plus"
        )

    def chat(self) -> Tuple[Response, int]:
        try:
            print(f"开始时间{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            data = request.get_json()
            # 新增单次问答选项处理
            singleTurn = data.get('singleTurn', False)
            dialogId = data.get('dialogId', str(uuid.uuid4()))
            model=self.models.get(data.get('model', 'qwen-plus'))
            # 根据模式构建消息
            mode = data.get('mode', '0')
            msg= data.get('message', '')
            messages = self._build_messages(mode, msg, dialogId, singleTurn)

            return self._get_response(messages, dialogId, data.get('stream', False), singleTurn,model)
        except Exception as e:
            return self._handle_error(e)
        
    def save(self) -> Tuple[Response, int]:
        if not (dialogId := request.json.get('dialogId')):
            return jsonify({"error": "需要dialogId"}), 400
        return jsonify(self._save_dialog(dialogId)), 200 if dialogId in self.active_dialogs else 404

    def load(self) -> Tuple[Response, int]:
        data = request.json
        dialogId = data.get('dialogId',str(uuid.uuid4))
        self.active_dialogs[dialogId] = [self.presets.get(data.get('mode', '0'), {}).copy()] if data.get('mode', '0') != '0' else []
        self.active_dialogs[dialogId].extend(data.get('messages', []))
        return jsonify({"dialogId": dialogId}), 200

    def _build_messages(self, mode: str, msg: str, dialogId: str, singleTurn: bool) -> List[Dict]:
        if singleTurn:
            # 单次问答模式直接构建消息
            messages = []
            if mode != '0' and (preset := self.presets.get(mode)):
                messages.append(preset.copy())
            messages.append({"role": "user", "content": msg})
            return messages
        else:
            # 原有上下文模式
            dialogId = self._manage_dialog(dialogId, mode)
            self._add_msg(dialogId, {"role": "user", "content": msg})
            return self.active_dialogs[dialogId]

    def _manage_dialog(self, dialogId: Optional[str], mode: str) -> str:
        if dialogId and dialogId in self.active_dialogs:
            return dialogId
        if mode not in self.presets:
            raise ValueError(f"无效模式，可选: {','.join(self.presets)}")
        new_id = dialogId or str(uuid.uuid4())
        self.active_dialogs[new_id] = [self.presets[mode].copy()] if self.presets[mode] else []
        return new_id

    def _add_msg(self, dialogId: str, msg: Dict):
        if not (isinstance(msg, dict) and msg.get("role") in ("user","assistant","system") and isinstance(msg.get("content"), str)):
            raise ValueError("消息格式错误")
        self.active_dialogs[dialogId].append(msg)
        # 检查对话记录数量，不包括系统提示
        non_system_messages = [m for m in self.active_dialogs[dialogId] if m["role"] != "system"]
        if len(non_system_messages) > 15:
            # 找到最早的非系统提示消息并删除
            for i, m in enumerate(self.active_dialogs[dialogId]):
                if m["role"] != "system":
                    del self.active_dialogs[dialogId][i]
                    break

    def _get_response(self, messages: List[Dict], dialogId: str, stream: bool, singleTurn: bool,model_config:AIModelConfig) -> Tuple[Response, int]:
        try:
            print(f"请求时间{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            resp = requests.post(
                model_config.api_base,  # ✅ 直接使用模型配置对象
                headers=model_config.get_auth_headers(),  # ✅ 直接使用认证头字典
                json=model_config.build_payload(messages, stream),
                stream=stream
                )
            resp.raise_for_status()
            print(f"响应时间{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            # 仅当非单次模式时保存对话记录
            if not singleTurn:
                return self._stream(resp, dialogId) if stream else self._normal(resp, dialogId)
            else:
                return self._stream(resp, None) if stream else self._normal(resp, None)
        except requests.RequestException as e:
            raise Exception(f"API请求失败: {str(e)}") from e

    # 修改流式响应处理
    def _stream(self, resp: requests.Response, dialogId: Optional[str]) -> Response:
        def generator():
            content = []
            for line in resp.iter_lines():
                if line and (data := line.decode().strip()).startswith('data:'):
                    try:
                        chunk = json.loads(data[5:])
                        if text := chunk.get('choices',[{}])[0].get('delta',{}).get('content'):
                            content.append(text)
                            yield f"data: {json.dumps({'content': text})}\n\n"
                    except json.JSONDecodeError: pass
            print(f"流式响应结束,时间{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", flush=True)
            # 仅在需要保存对话时记录
            if dialogId:
                self._add_msg(dialogId, {"role": "assistant", "content": ''.join(content)})
        return Response(generator(), mimetype='text/event-stream')

    # 修改普通响应处理
    def _normal(self, resp: requests.Response, dialogId: Optional[str]) -> Tuple[Response, int]:
        content = resp.json()['choices'][0]['message']['content']
        if dialogId:
            self._add_msg(dialogId, {"role": "assistant", "content": content})
            return jsonify({"dialogId": dialogId, "response": content}), 200
        else:
            return jsonify({"response": content}), 200

    def _save_dialog(self, dialogId: str) -> Dict:
        return {"dialogId": dialogId, "messages": self.active_dialogs.get(dialogId, [])}

    def _handle_error(self, error: Exception) -> Tuple[Response, int]:
        code = 500
        if isinstance(error, ValueError): code = 400
        return jsonify({"error": str(error)}), code